package com.dyp.multi_data_source.config;

import com.zaxxer.hikari.HikariDataSource;
import jakarta.persistence.EntityManagerFactory;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.boot.orm.jpa.EntityManagerFactoryBuilder;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.jpa.repository.config.EnableJpaRepositories;
import org.springframework.orm.jpa.JpaTransactionManager;
import org.springframework.orm.jpa.JpaVendorAdapter;
import org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean;
import org.springframework.orm.jpa.vendor.Database;
import org.springframework.orm.jpa.vendor.HibernateJpaVendorAdapter;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.annotation.EnableTransactionManagement;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;
import java.util.function.Function;

@Configuration
@EnableJpaRepositories(basePackages = "com.dyp.multi_data_source.repository")
@EnableTransactionManagement
public class MultiDataSourceConfig {

    @Bean(name = "primaryDataSource")
    @ConfigurationProperties(prefix = "spring.datasource.primary")
    public DataSource masterDataSource() {
        return DataSourceBuilder.create().build();
    }

    @Bean(name = "secondaryDataSource")
    @ConfigurationProperties(prefix = "spring.datasource.secondary")
    public DataSource slaveDataSource() {
        return DataSourceBuilder.create().build();
    }

    @Bean
    public DataSource dataSource(@Qualifier("primaryDataSource") DataSource master,
                                 @Qualifier("secondaryDataSource") DataSource slave) {
        DynamicDataSource dynamicDataSource = new DynamicDataSource();
        Map<Object, Object> targetDataSources = new HashMap<>();
        targetDataSources.put("master",  master);
        targetDataSources.put("slave",  slave);
        dynamicDataSource.setTargetDataSources(targetDataSources);
        dynamicDataSource.setDefaultTargetDataSource(master);
        return dynamicDataSource;
    }

    // 显式定义EntityManagerFactoryBuilder Bean  3.3.4之前的版本
//    @Bean
//    public EntityManagerFactoryBuilder entityManagerFactoryBuilder(JpaVendorAdapter jpaVendorAdapter) {
//        return new EntityManagerFactoryBuilder(jpaVendorAdapter, new HashMap<>(), null);
//    }
    @Bean
    public Function<DataSource, Map<String, ?>> jpaPropertiesFactory() {
        return dataSource -> {
            Map<String, Object> properties = new HashMap<>();
            properties.put("hibernate.hbm2ddl.auto",  "update");
            properties.put("hibernate.format_sql",  true);
            // 可以根据数据源动态设置属性
            if (dataSource instanceof HikariDataSource) {
                properties.put("hibernate.connection.provider_class",
                        "com.zaxxer.hikari.hibernate.HikariConnectionProvider");
            }
            return properties;
        };
    }
    @Bean
    public EntityManagerFactoryBuilder entityManagerFactoryBuilder(
            JpaVendorAdapter jpaVendorAdapter,
            Function<DataSource, Map<String, ?>> jpaPropertiesFactory) {
        // PersistenceUnitManager 可以是null（除非需要自定义PersistenceUnit）
        return new EntityManagerFactoryBuilder(
                jpaVendorAdapter,
                jpaPropertiesFactory,
                null  // PersistenceUnitManager
        );
    }

    // 定义JpaVendorAdapter
    @Bean
    public JpaVendorAdapter jpaVendorAdapter() {
        HibernateJpaVendorAdapter adapter = new HibernateJpaVendorAdapter();
        adapter.setDatabase(Database.MYSQL);
        adapter.setShowSql(true);
        adapter.setGenerateDdl(true);
        return adapter;
    }


    @Bean(name = "entityManagerFactory")
    public LocalContainerEntityManagerFactoryBean secondaryEntityManagerFactory(
            EntityManagerFactoryBuilder builder,
            DataSource dataSource) {
        return builder
                .dataSource(dataSource)
                .packages("com.dyp.multi_data_source.entity")
                .build();
    }

    @Bean(name = "transactionManager")
    public PlatformTransactionManager secondaryTransactionManager(
            @Qualifier("entityManagerFactory") EntityManagerFactory entityManagerFactory) {
        return new JpaTransactionManager(entityManagerFactory);
    }
}
